import pandas as pd
import pymysql
from datetime import datetime
import os

def get_book_info(filename):
    # 从文件名解析学段、年级和学期信息
    stage = ''
    grade = ''
    term = ''
    
    if '高中' in filename:
        stage = '高中'
    elif '初中' in filename:
        stage = '初中'
    else:
        stage = '小学'
    
    if '三年级' in filename:
        grade = '三年级'
    elif '四年级' in filename:
        grade = '四年级'
    elif '五年级' in filename:
        grade = '五年级'
    elif '六年级' in filename:
        grade = '六年级'
    elif '初中第一册' in filename:
        grade = '初一'
    elif '高中必修一' in filename:
        grade = '高一'
    
    if '上册' in filename:
        term = '上册'
    elif '下册' in filename:
        term = '下册'
    
    return stage, grade, term

try:
    # 连接数据库
    print('正在连接数据库...')
    conn = pymysql.connect(
        host='localhost',
        user='root',
        password='root',
        database='ry-vue',
        charset='utf8mb4'
    )
    print('数据库连接成功！')

    # 清空现有数据
    print('正在清空现有数据...')
    with conn.cursor() as cursor:
        cursor.execute('TRUNCATE TABLE sys_word_book')
    print('现有数据已清空！')

    # 获取ori-data目录下的所有Excel文件
    excel_files = ['英语单词表_格式模板(2).xlsx']
    total_count = 0

    # 插入数据
    with conn.cursor() as cursor:
        for excel_file in excel_files:
            print(f'\n正在处理文件：{excel_file}')
            # 读取Excel文件
            df = pd.read_excel(os.path.join('ori-data', excel_file))
            print('Excel文件读取成功，开始处理数据...')
            print('列名：', df.columns.tolist())
            print('数据行数：', len(df))

            # 获取书本信息
            stage, grade, term = get_book_info(excel_file)
            book_name = os.path.splitext(excel_file)[0]

            count = 0
            for index, row in df.iterrows():
                sql = '''INSERT INTO sys_word_book 
                         (word, phonetic, syllable, part_of_speech, chinese_meaning, example_sentence,
                          example_translate, book_name, stage, grade, term, 
                          create_by, create_time, status) 
                         VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, 'admin', %s, '0')'''
                
                # 获取每列的值，如果列不存在则使用None
                word = str(row['单词']) if not pd.isna(row['单词']) else None
                phonetic = str(row['单词音标']) if '单词音标' in df.columns and not pd.isna(row['单词音标']) else None
                syllable = str(row['音节']) if '音节' in df.columns and not pd.isna(row['音节']) else None
                part_of_speech = str(row['词性']) if '词性' in df.columns and not pd.isna(row['词性']) else None
                chinese_meaning = str(row['单词含义']) if '单词含义' in df.columns and not pd.isna(row['单词含义']) else None
                example_sentence = str(row['例句']) if '例句' in df.columns and not pd.isna(row['例句']) else None
                example_translate = str(row['例句翻译']) if '例句翻译' in df.columns and not pd.isna(row['例句翻译']) else None
                
                cursor.execute(sql, (
                    word,
                    phonetic,
                    syllable,
                    part_of_speech,
                    chinese_meaning,
                    example_sentence,
                    example_translate,
                    book_name,
                    stage,
                    grade,
                    term,
                    datetime.now()
                ))
                count += 1
                if count % 100 == 0:
                    print(f'已处理 {count} 条数据')
                    conn.commit()
            
            conn.commit()
            total_count += count
            print(f'{excel_file} 导入完成！共导入 {count} 条数据')

    print(f'\n所有文件处理完成！总共导入 {total_count} 条数据')

except Exception as e:
    print(f'错误：{str(e)}')
    if 'conn' in locals():
        conn.rollback()
finally:
    if 'conn' in locals():
        conn.close()
        print('数据库连接已关闭')